package cn.sc.summer.token.encrypt;

import cn.sc.summer.constant.encrypt.EncryptKeyConstant;
import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.*;

/**
 * 类名：RSA非对称加密类
 *
 * @author a-xin
 * @date 2024/2/5 15:32
 */
@Slf4j
public class RSAUtil {

    /**
     * 字符集
     */
    private static final String CHARSET = "UTF-8";

    /**
     * 加密算法
     */
    private final static String ALGORITHM_RSA = "RSA";

    private static final Cipher CIPHER;

    static {
        try {
            CIPHER = Cipher.getInstance(ALGORITHM_RSA);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 直接生成公钥、私钥对象
     *
     * @param modulus 模长
     * @return 公钥私钥对象
     * @throws NoSuchAlgorithmException 异常
     */
    public static List<Key> getRSAKeyObject(int modulus) throws NoSuchAlgorithmException {
        List<Key> keyList = new ArrayList<>(2);
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(ALGORITHM_RSA);
        keyPairGen.initialize(modulus);
        KeyPair keyPair = keyPairGen.generateKeyPair();
        keyList.add(keyPair.getPublic());
        keyList.add(keyPair.getPrivate());
        return keyList;
    }

    /**
     * 生成公钥、私钥的字符串
     *
     * @param modulus 模长
     * @return 公钥私钥
     * @throws NoSuchAlgorithmException 异常
     */
    public static Map<String, String> getRSAKeyString(int modulus) throws NoSuchAlgorithmException {
        Map<String, String> keyPairMap = new HashMap<>();
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(ALGORITHM_RSA);
        keyPairGen.initialize(modulus);
        KeyPair keyPair = keyPairGen.generateKeyPair();
        String publicKey = Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
        String privateKey = Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
        keyPairMap.put(EncryptKeyConstant.PUBLIC_KEY, publicKey);
        keyPairMap.put(EncryptKeyConstant.PRIVATE_KEY, privateKey);
        return keyPairMap;
    }

    /**
     * Java中RSAPublicKeySpec、X509EncodedKeySpec支持生成RSA公钥，此处使用X509EncodedKeySpec生成
     *
     * @param publicKey 公钥
     * @return 公钥对象
     * @throws Exception 异常
     */
    public static RSAPublicKey getPublicKey(String publicKey) throws Exception {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
        byte[] keyBytes = Base64.getMimeDecoder().decode(publicKey);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
        return (RSAPublicKey) keyFactory.generatePublic(spec);
    }

    /**
     * Java中只有RSAPrivateKeySpec、PKCS8EncodedKeySpec支持生成RSA私钥，此处使用PKCS8EncodedKeySpec生成
     *
     * @param privateKey 私钥
     * @return 私钥对象
     * @throws Exception 异常
     */
    public static RSAPrivateKey getPrivateKey(String privateKey) throws Exception {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
        byte[] keyBytes = Base64.getMimeDecoder().decode(privateKey);
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
        return (RSAPrivateKey) keyFactory.generatePrivate(spec);
    }

    /**
     * 公钥加密
     *
     * @param data      加密数据
     * @param publicKey 公钥对象
     * @return 加密后字符串
     * @throws Exception 异常
     */
    public static String encryptByPublicKey(String data, RSAPublicKey publicKey) throws Exception {
        CIPHER.init(Cipher.ENCRYPT_MODE, publicKey);
        int modulusSize = publicKey.getModulus().bitLength() / 8;
        int maxSingleSize = modulusSize - 11;
        byte[][] dataArray = splitArray(data.getBytes(CHARSET), maxSingleSize);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        for (byte[] s : dataArray) {
            out.write(CIPHER.doFinal(s));
        }
        return Base64.getEncoder().encodeToString(out.toByteArray());
    }

    /**
     * 公钥解密
     *
     * @param data      解密数据
     * @param publicKey 公钥
     * @return 解密后数据
     * @throws Exception 异常
     */
    public static String decryptByPublicKey(String data, RSAPublicKey publicKey) throws Exception {
        CIPHER.init(Cipher.DECRYPT_MODE, publicKey);
        int modulusSize = publicKey.getModulus().bitLength() / 8;
        byte[] dataBytes = data.getBytes(CHARSET);
        byte[] decodeData = Base64.getMimeDecoder().decode(dataBytes);
        byte[][] splitArrays = splitArray(decodeData, modulusSize);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        for (byte[] arr : splitArrays) {
            out.write(CIPHER.doFinal(arr));
        }
        return out.toString();
    }

    /**
     * 私钥解密
     *
     * @param data       解密数据
     * @param privateKey 私钥
     * @return 解密后数据
     * @throws Exception 异常
     */
    public static String decryptByPrivateKey(String data, RSAPrivateKey privateKey) throws Exception {
        CIPHER.init(Cipher.DECRYPT_MODE, privateKey);
        int modulusSize = privateKey.getModulus().bitLength() / 8;
        byte[] dataBytes = data.getBytes(CHARSET);
        byte[] decodeData = Base64.getMimeDecoder().decode(dataBytes);
        byte[][] splitArrays = splitArray(decodeData, modulusSize);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        for (byte[] arr : splitArrays) {
            out.write(CIPHER.doFinal(arr));
        }
        return out.toString();
    }

    /**
     * 私钥加密
     *
     * @param data       加密数据
     * @param privateKey 私钥
     * @return 加密后数据
     * @throws Exception 异常
     */
    public static String encryptByPrivateKey(String data, RSAPrivateKey privateKey) throws Exception {
        CIPHER.init(Cipher.ENCRYPT_MODE, privateKey);
        int modulusSize = privateKey.getModulus().bitLength() / 8;
        int maxSingleSize = modulusSize - 11;
        byte[][] dataArray = splitArray(data.getBytes(CHARSET), maxSingleSize);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        for (byte[] s : dataArray) {
            out.write(CIPHER.doFinal(s));
        }
        return Base64.getMimeEncoder().encodeToString(out.toByteArray());
    }

    /**
     * 按指定长度切分数组
     *
     * @param data 数组
     * @param len  长度
     * @return 分割后数组
     */
    private static byte[][] splitArray(byte[] data, int len) {
        int dataLen = data.length;
        if (dataLen <= len) {
            return new byte[][]{data};
        }
        byte[][] result = new byte[(dataLen - 1) / len + 1][];
        int resultLen = result.length;
        for (int i = 0; i < resultLen; i++) {
            if (i == resultLen - 1) {
                int slen = dataLen - len * i;
                byte[] single = new byte[slen];
                System.arraycopy(data, len * i, single, 0, slen);
                result[i] = single;
                break;
            }
            byte[] single = new byte[len];
            System.arraycopy(data, len * i, single, 0, len);
            result[i] = single;
        }
        return result;
    }

}
